package com.example.emos.wx.utils;

import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HtmlUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;
import java.nio.charset.Charset;
import java.util.*;

public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    private String body;
//    private byte[] body;

    private HashMap<String, String[]> params = new HashMap<>();

    public XssHttpServletRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        params.putAll(request.getParameterMap());
    }

    public XssHttpServletRequestWrapper(HttpServletRequest request, Map<String, Object> extendParams) throws IOException {
        this(request);
        addAllParameters(extendParams);
    }

    @Override
    public String getParameter(String name) {
        String[] values = params.get(name);
        if (values == null || values.length == 0) {
            return null;
        }
        String val = values[0];
        if (!StrUtil.isEmpty(val)) {
            val = HtmlUtil.filter(val);
        }
        return val;
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] values = params.get(name);
        if (values != null) {
            for (int i = 0; i < values.length; i++) {
                String value = values[i];
                if (!StrUtil.isEmpty(value)) {
                    value = HtmlUtil.filter(value);
                }
                values[i] = value;
            }
        }
        return values;
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        for (Map.Entry<String, String[]> stringEntry : params.entrySet()) {
            String[] values = stringEntry.getValue();
            for (int i = 0; i < values.length; i++) {
                String value = values[i];
                if (!StrUtil.hasEmpty(value)) {
                    values[i] = HtmlUtil.filter(value);
                }
            }
        }
        return Collections.unmodifiableMap(params);
    }

    @Override
    public Enumeration<String> getParameterNames() {
        return Collections.enumeration(params.keySet());
    }

    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(getInputStream()));
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        body = read(super.getInputStream());
        ByteArrayInputStream inputStream = new ByteArrayInputStream(body.getBytes());
        return new ServletInputStream() {
            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setReadListener(ReadListener readListener) {

            }

            @Override
            public int read() throws IOException {
                return inputStream.read();
            }
        };
    }

    /**
     * 读取请求体并转换为字节数组
     * 方案1
     *
     * @param input
     * @return
     * @throws IOException
     */
    private String read(InputStream input) throws IOException {
        InputStreamReader reader = new InputStreamReader(input, Charset.forName("UTF-8"));
        BufferedReader buffer = new BufferedReader(reader);
        StringBuffer bodyString = new StringBuffer();
        String line = buffer.readLine();
        while (line != null) {
            bodyString.append(line);
            line = buffer.readLine();
        }
        buffer.close();
        reader.close();
        input.close();
        Map<String, Object> map = JSONUtil.parseObj(bodyString);
        HashMap<String, Object> result = new HashMap<>(map.size());
        for (String key : map.keySet()) {
            Object val = map.get(key);
            if (val instanceof String) {
                if (!StrUtil.isEmpty(val.toString())) {
                    result.put(key, HtmlUtil.filter(val.toString()));
                }
            } else {
                result.put(key, val);
            }
        }
        String str = JSONUtil.toJsonStr(result);

        return str;

    }

    /**
     * 读取请求体并转换为字节数组
     * 方案2
     *
     * @param input
     * @return
     * @throws IOException
     */
    private byte[] read2(InputStream input) throws IOException {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        byte[] buffer = new byte[1024];

        int len = 0;
        while ((len = input.read(buffer)) != -1) {
            outputStream.write(buffer, 0, len);
        }
        input.close();

        return outputStream.toByteArray();

    }

    /**
     * 批量添加参数
     *
     * @param otherParams
     */
    private void addAllParameters(Map<String, Object> otherParams) {
        Set<Map.Entry<String, Object>> entries = otherParams.entrySet();
        for (Map.Entry<String, Object> entry : entries) {
            addParameters(entry.getKey(), entry.getValue());
        }
    }

    /**
     * 添加单个参数
     *
     * @param key
     * @param values
     */
    private void addParameters(String key, Object values) {
        // 判断是否已经存在
        String[] parameterValues = getParameterValues(key);
        if (parameterValues == null) {
            if (values != null) {
                if (values instanceof String) {
                    params.put(key, new String[]{(String) values});
                } else if (values instanceof String[]) {
                    params.put(key, (String[]) values);
                } else {
                    params.put(key, new String[]{String.valueOf(values)});
                }
            }
        }

    }
}
